

import os
import time
import h5py as h5
import numpy as np
import utils as ut
import pandas as pd

from tqdm import tqdm
from pyflann import *


def test_FLANN(dataset, test_set, test_truth, target_precision):

    flann = FLANN()
    params = flann.build_index(dataset, algorithm='autotuned', target_precision=target_precision, log_level="error")

    sum_time = 0
    sum_recall = 0
    for _ in range(10):
        start = time.perf_counter()
        for query in test_set:
            result, dists = flann.nn_index(query, checks=params["checks"], num_neighbors=10)
        sum_time += time.perf_counter() - start

        for i, query in enumerate(test_set):
            result, dists = flann.nn_index(query, checks=params["checks"], num_neighbors=10)
            truth = test_truth[i][0][:10]
            sum_recall += ut.calc_recall(truth, result[0])

    recall = sum_recall / (10 * len(test_set))
    query_time = sum_time / (10 * len(test_set))

    return query_time, recall


def generate_form(data, columns, file_name):

    res = pd.DataFrame(data, columns=columns)

    if not os.path.exists(file_name):
        res.to_csv(file_name, index=False)
    else:
        res.to_csv(file_name, index=False, header=False, mode='a')


def main(dir_path, target_dir='./dataset', from_name='res.csv'):
    precision_dict = {'audio': 0.95, 'sun': 0.91,
                      'enron': 0.96, 'nuswide': 0.77,
                      'notre': 0.97}


    for file_name in os.listdir(dir_path):
        ut.generate_dataset(os.path.join(dir_path, file_name), target_dir)

    file_name_list = []
    query_time_list = []
    recall_list = []

    for file_name in os.listdir(target_dir):
        dataset, test_set, test_truth = ut.load_datasets(os.path.join(target_dir, file_name), test=True)

        file_name, extension = os.path.splitext(file_name)
        query_time, recall = test_FLANN(dataset, test_set, test_truth, precision_dict[file_name])

        file_name_list.append(file_name)
        query_time_list.append(query_time)
        recall_list.append(recall)

    data = {
        'Dataset': file_name_list,
        'queryTime': query_time_list,
        'Recall': recall_list
    }

    columns = ['Dataset', 'queryTime', 'Recall']

    generate_form(data, columns, from_name)




if __name__ == '__main__':
    root_dir = 'E:'
    dir_path = os.path.join(root_dir, 'target_datasets')
    main(dir_path)


    pass